import numpy as np

# 查找数组中值为 val 的元素的下标，如果找不到，则返回离其最近的上一个值的下标
def get_index(arr, val):
    """
    查找数组中值为 val 的元素的下标，如果找不到，则返回离其最近的上一个值的下标
    :param arr: 深度数组
    :param val: 要查找的数据
    :return: 下标
    """
    indices = np.where(arr == val)[0]
    if indices.size > 0:
        # 如果找到了 deep，获取其第一个下标
        index = indices[0]
    else:
        # 如果没找到 deep，向下取最近的数据的下标
        diff = arr - val
        valid_indices = np.where(diff < 0)[0]
        if valid_indices.size > 0:
            index = valid_indices[diff[valid_indices].argmax()]
        else:
            # 如果所有元素都大于 val，取第一个元素的下标
            index = 0

    return index

# 绘制拉平数据，分层
def wiggle_plot(file, cfg, ax, layer):
    """
    绘制拉平数据，分层
    :param file: 拉平数据
    :param cfg: 参数数组
    :param ax: 坐标系
    :param layer: 分层数据
    :return: [fbmin, fbmax, ax, fctime]，最小初至，最大初至，坐标系，层深对应的初至
    """
    perc = 99.0
    data = file["data"]
    rgba = list(cfg['colour']) + [cfg['opacity']]
    sc = np.percentile(data, perc)  # Normalization factor，计算数据集中特定百分位数的值，相当于去除异常数据后的最大值
    stp = int(cfg['trange'][0] * 1000.0 / file["dt"])  # 计算时间点索引，开始
    etp = int(cfg['trange'][1] * 1000.0 / file["dt"])  # 计算时间点索引，结束
    wigdata = data[::cfg['skip'], stp: etp: cfg['skip2']]  # 对道和时间进行抽稀，并对时间进行截取
    brfk = 1  # 跳过绘制的道数+1
    file["cdp"] = file["cdp"] + cfg["kb"]  # 深度加补芯高
    xpos = np.array(file["cdp"])  # 转为numpy数组，深度
    brfp = 0  # 跳过绘制的道数
    fbmax = 0.0  # 最大初至时间
    fbmin = 8000.0  # 最小初始时间

    # 绘制拉平
    # 对要绘制的深度范围进行筛选，前后都加补芯高，相当于不加补芯高
    # 计算最大和最小初至时间
    # 将道数据绘制到对应深度的位置
    # brfjj为循环自变量，x为深度，trace为道数据
    for brfjj, x, trace in zip(range(0, len(xpos)), xpos, wigdata):
        # Compute high resolution trace.
        # 先筛选后进行下面的操作，所以brfk为开始绘制的第2道的索引
        if x > cfg["lped"] + cfg["kb"] or x < cfg["lpsd"] + cfg["kb"]:
            brfp = brfp + 1  # 跳过绘制的道数
            brfk = brfk + 1  # 跳过绘制的道数+1
            print('跳过第', brfp, '道，共', len(xpos), '道')
            continue

        if file["fb"][brfjj] > fbmax:
            fbmax = file["fb"][brfjj]
        if file["fb"][brfjj] < fbmin:
            fbmin = file["fb"][brfjj]
        if brfk == xpos.size:  # 当跳过绘制的道数+1为总道数时，此时还有1道没有跳过
            brfk = brfk - 1
        # 深度 - 振幅大小 * 第1道和第2道深度差 * 道数据(0~1)
        # 将道数据与深度结合
        if cfg['direction'] == 0:
            amp = x - cfg['lpgain'] * (xpos[brfk] - xpos[brfk - 1]) * trace / sc
        else:
            amp = x - cfg['lpgain'] * (xpos[brfk - 1] - xpos[brfk]) * trace / sc

        # 对道数据进行插值
        # 时间数组，起始，结束，步长
        t = np.arange(cfg['trange'][0] * 1000, cfg['trange'][1] * 1000, file["dt"] * cfg["skip2"])
        # 创建t[0]和t[-1]之间的等间距样本数组，数据个数为10 * timeArr.size，包括t[-1]
        # 相当于将数组timeArr的颗粒度由1变为0.1
        hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
        # 利用线性插值公式得到新的amp数组，y = y1 + (y2 - y1) * ((x - x1) / (x2 - x1))
        # 相当于使用timeArr和newTrace计算出一个拟合曲线，然后计算出hypertime对应的data中的地震数据trace
        hyperamp = np.interp(hypertime, t, amp)  # 差值后的道数据

        # Plot the line, then the fill.
        for brfi in range(len(hyperamp)):
            # 初至之前的数据不绘制
            if hypertime[brfi] < file["fb"][brfp]:
                hyperamp[brfi] = x
            # 走廊切除深度之前的数据，走廊视窗外的部分不显示
            if x < cfg["stkdep"]:
                if hypertime[brfi] > file["fb"][brfp] + cfg['stkw']:
                    hyperamp[brfi] = x

        # 绘制该道数据
        ax.plot(hyperamp, hypertime, 'k', lw=1.2 * cfg['lineweight'])
        # 对波形进行填充
        # hypertime,  # 填充区域的y轴坐标
        # hyperNewTrace,  # 填充区域的下边界的x轴坐标
        # x,  # 填充区域的上边界的x轴坐标
        # where=hyperNewTrace < x,  # 指定哪些区域需要填充
        # facecolor=rgba,  # 填充颜色和透明度
        # lw=0,  # 设置填充区域的边界线的宽度为 0，意味着不显示边界线
        ax.fill_betweenx(hypertime, hyperamp, x, where=hyperamp < x, facecolor=rgba, lw=0,)

        #  更新绘制道的下标
        brfk = brfk + 1
        brfp = brfp + 1

    # 沿初至绘制红线
    ax.plot(file["cdp"], file["fb"], 'r', lw=2.5 * cfg['lineweight'])
    # 绘制蓝线
    # 曲线部分
    index = get_index(file["cdp"], cfg["stkdep"])
    line_cdp = file["cdp"][:index + 1]
    line_fb = file["fb"][:index + 1] + cfg["stkw"]
    ax.plot(line_cdp, line_fb, 'b', lw=2.5 * cfg['lineweight'])
    # 直线部分
    index2 = get_index(hypertime, file["fb"][index] + cfg["stkw"])
    line_time = hypertime[index2:]
    line_deep = np.full(line_time.size, file["cdp"][index])
    ax.plot(line_deep, line_time, 'b', lw=2.5 * cfg['lineweight'])

    # 设置坐标走
    ax.invert_yaxis()
    ax.set_xlim(cfg["lpsd"] + cfg["kb"] - 20, cfg["lped"] + cfg["kb"] + 20)
    ax.invert_xaxis()

    # 绘制分层
    brfpp = 0  # 分层数据的下标
    fctime = []  # 存放分层数据对应的初至
    dep = layer["dep"]  # 底深
    ywdhpy = layer["ywdhpy"]  # 代号深度
    ywdh = layer["ywdh"]  # 英文代号
    font1 = layer["font1"]  # 代号大小
    font = layer["font"]  # 文字大小
    py = layer["py"]  # 文字深度偏移
    fzt = layer["fzt"]  # 文字时间
    zu = layer["zu"]  # 分组
    if cfg["fcx"] == 1:
        fctime = np.arange(len(dep))

        # 计算分层深度对应的初至时间
        for ii in dep:
            for jj in range(1, len(file["cdp"])):
                if ii == file["cdp"][jj]:  # 如果cdp中有分层深度，则记录该深度对应的初至时间
                    fctime[brfpp] = file["fb"][jj]
                    brfpp = brfpp + 1
                elif ((ii > file["cdp"][jj - 1] and ii < file["cdp"][jj]) or (
                        ii < file["cdp"][jj - 1] and ii > file["cdp"][jj])):  # 如果没有，则计算该深度对应的初至时间并记录
                    fctime[brfpp] = (ii - file["cdp"][jj]) * (
                            file["fb"][jj - 1] - file["fb"][jj]) / (
                                        (file["cdp"][jj - 1] - file["cdp"][jj])) + file["fb"][jj]
                    brfpp = brfpp + 1

        for ii in range(0, len(dep)):
            # 跳过超出绘制深度范围的分层
            if dep[ii] > cfg["lped"] + cfg["kb"] or dep[ii] < cfg["lpsd"] + cfg["kb"]:
                continue

            # 绘制分层
            ax.plot([dep[ii], dep[ii], 0], [0, fctime[ii], fctime[ii]], 'r', lw=2.5 * cfg['lineweight'])

            # 绘制英文代号
            ax.text(ywdhpy[ii], fctime[ii], ywdh[ii],
                    transform=ax.transData,
                    verticalalignment='center',  # 垂直居中
                    horizontalalignment='center',  # 水平居中
                    fontsize=font1[ii],
                    rotation=270, color='k')

            # 绘制分组
            if ii >= 1 and (ii < (len(dep) - 1)):
                if (dep[ii] - dep[ii - 1] < font[ii] * 16) and (
                        (dep[ii + 1] - dep[ii] < font[ii] * 16)):
                    ax.text(dep[ii] + font[ii] * 4 + py[ii],
                            fzt[ii], zu[ii], transform=ax.transData,
                            fontsize=font[ii],
                            rotation=270, color='r')
                elif (dep[ii] - dep[ii - 1] < font[ii] * 16) and (
                        (dep[ii + 1] - dep[ii] > font[ii] * 16)):
                    ax.text(dep[ii] + font[ii] * 8 + py[ii],
                            + fzt[ii], zu[ii], transform=ax.transData,
                            fontsize=font[ii],
                            rotation=270, color='r')
                elif (dep[ii] - dep[ii - 1] > font[ii] * 16) and (
                        (dep[ii + 1] - dep[ii] < font[ii] * 16)):
                    ax.text(dep[ii], fzt[ii] + py[ii], zu[ii],
                            transform=ax.transData,
                            fontsize=font[ii],
                            rotation=270, color='r')
                elif (dep[ii] - dep[ii - 1] > font[ii] * 16) and (
                        (dep[ii + 1] - dep[ii] > font[ii] * 16)):
                    ax.text(dep[ii] + font[ii] * 8 + py[ii], fzt[ii],
                            zu[ii], transform=ax.transData,
                            fontsize=font[ii],
                            rotation=270, color='r')
            else:
                ax.text(dep[ii] + py[ii], fzt[ii], zu[ii],
                        transform=ax.transData, fontsize=font[ii],
                        rotation=270, color='r')


    return [fbmin, fbmax, ax, fctime]

# 用于绘制走廊和滤波走廊
def wiggle_plot1(file, cfg, ax, gain, fctime, fbmax, fbmin, stack, times=0):
    """
    用于绘制走廊和滤波走廊
    :param file: 走廊数据
    :param cfg: 参数对象
    :param ax: 坐标系对象
    :param gain: 振幅大小
    :param fctime: 分层深度对应的初至数组
    :param fbmax: 初至最大值
    :param fbmin: 初至最小值
    :param times: 时移，默认为0
    :return: 坐标系对象
    """
    trace_times = cfg["times"]  # 对走廊和滤波走廊数据，按照斜井裁剪时使用
    perc = 99.0
    data = file["data"]
    rgba = list(cfg['colour']) + [cfg['opacity']]
    sc = np.percentile(data, perc)  # Normalization factor，计算数据集中特定百分位数的值，相当于去除异常数据后的最大值
    stp = int((cfg['trange'][0] * 1000.0 + times) / file["dt"])  # 计算时间点索引，开始，加时移times
    etp = int((cfg['trange'][1] * 1000.0 + times) / file["dt"])  # 计算时间点索引，结束，加时移times
    wigdata = data[::cfg["skip"], stp: etp: cfg["skip2"]]  # 对道和时间进行抽稀，并对时间进行截取

    # 设置背景透明
    ax.set_facecolor((1, 1, 1, 0))  # RGBA 颜色，其中 A（透明度）为 0 表示完全透明

    # 将斜井数据部分背景设置为白色
    border_trace = []
    border_time = []
    trace_num = stack["trace_num"] + 1  # 这里的道数为走廊道数+1
    for i in range(len(stack["trace_offset"])):
        border_trace.append(stack["trace_offset"][i]["trace_index"] * 20)  # 道的y轴，左侧第一道为13，右侧第一道为11
        border_time.append(stack["trace_offset"][i]["non_zero_time"] - trace_times)  # 该道为0的时间点，即每道井斜偏移的时间点
    hyper_time = np.linspace(border_time[0], border_time[-1], (10 * len(border_time) - 1) + 1)
    hyper_trace = np.interp(hyper_time, border_time, border_trace)  # 道的y轴，左侧第一道为13，右侧第一道为11
    # 添加白色背景
    if stack["flag"] == -1:  # 井斜左偏
        ax.fill_betweenx([cfg['trange'][0] * 1000, hyper_time[0]], [hyper_trace[0], hyper_trace[0]],
                         [hyper_trace[0] + trace_num * 20, hyper_trace[0] + trace_num * 20], color='white', alpha=1.0,
                         interpolate=True)  # 直线部分
        ax.fill_betweenx(hyper_time, hyper_trace, hyper_trace + trace_num * 20, color='white', alpha=1.0,
                         interpolate=True)  # 弯曲部分
        ax.fill_betweenx([hyper_time[-1], cfg['trange'][1] * 1000], [0, 0], [trace_num * 20, trace_num * 20],
                         color='white', alpha=1.0, interpolate=True)  # 直线部分
    else:  # 井斜右偏
        ax.fill_betweenx([cfg['trange'][0] * 1000, hyper_time[0]], [0, 0],
                         [trace_num * 20, trace_num * 20], color='white', alpha=1.0,
                         interpolate=True)  # 直线部分
        ax.fill_betweenx(hyper_time, hyper_trace, hyper_trace - trace_num * 20, color='white', alpha=1.0,
                         interpolate=True)  # 弯曲部分
        ax.fill_betweenx([hyper_time[-1], cfg['trange'][1] * 1000], [hyper_trace[-1] - trace_num * 20 + 20, hyper_trace[-1] - trace_num * 20 + 20],
                         [hyper_trace[-1] + 20, hyper_trace[-1] + 20], color='white', alpha=1.0,
                         interpolate=True)  # 直线部分

    # 绘制波形前的预处理
    xpos = np.arange(20, (stack["trace_num_offset"] + 1) * 20, 20)  # 生成[20, 460]的数组，步长为20。默认每道的深度为20
    for x, trace in zip(xpos, wigdata):
        t = np.arange(cfg['trange'][0] * 1000, cfg['trange'][1] * 1000, file["dt"] * cfg["skip2"])
        # 过滤掉值为 0 的数据
        trace_index = int(x / 20)
        if stack["flag"] == -1:
            trace_offset_reverse = stack["trace_offset"][::-1]  # 数组翻转
            if trace_index <= stack["trace_start_offset"]:  # 左下方，13
                index = get_index(t, trace_offset_reverse[trace_index - 1]["non_zero_time"] - trace_times)
                trace = trace[index:]
                t = t[index:]
            if trace_index > stack["trace_num"] + 1:  # 右上方，12
                index = get_index(t, trace_offset_reverse[trace_index - stack["trace_num"] - 2]["non_zero_time"] - trace_times)
                trace = trace[:index+1]
                t = t[:index+1]
            if trace_index == stack["trace_num"] + 1:  # 第11道数据
                index = get_index(t, trace_offset_reverse[0]["non_zero_time"] - trace_times)
                trace = trace[:index + 1]
                t = t[:index + 1]
        else:  # 如是右侧偏移，则仅去除右上方的值为0的数据
            if trace_index > stack["trace_num"]:  # 右下方，11
                index = get_index(t, stack["trace_offset"][trace_index - stack["trace_num"] - 1]["non_zero_time"] - trace_times)
                trace = trace[index:]
                t = t[index:]
            if trace_index < stack["trace_num_offset"] - stack["trace_num"]: # 左上方，13
                index = get_index(t, stack["trace_offset"][trace_index]["non_zero_time"] - trace_times)
                trace_line = [x] * len(trace[index:])  # 直线
                trace_time = t[index:]  # 直线
                ax.plot(trace_line, trace_time, 'k', lw=1.2 * cfg['lineweight'])
                trace = trace[:index + 1]
                t = t[:index + 1]
            if trace_index == stack["trace_num_offset"] - stack["trace_num"]:  # 第13道数据
                index = get_index(t, stack["trace_offset"][-1]["non_zero_time"] - trace_times)
                trace = trace[:index + 1]
                t = t[:index + 1]

        # Compute high resolution trace.
        # 将道数据与深度结合
        amp = x - gain * 20 * trace / sc

        # 对道数据进行插值
        # 相当于将数组timerArr的颗粒度由1变为0.1
        hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
        # 利用下行差值公式得到新的amp数组
        hyperamp = np.interp(hypertime, t, amp)

        # 绘制该道数据
        ax.plot(hyperamp, hypertime, 'k', lw=1.2 * cfg['lineweight'])
        # 对波形进行填充
        ax.fill_betweenx(hypertime, hyperamp, x,
                         where=hyperamp > x,
                         facecolor=rgba,
                         lw=0,
                         )

    # 绘制两条曲线
    if stack["flag"] == -1:  # 向左偏
        # 第一条曲线
        ax.plot([hyper_trace[0], hyper_trace[0]], [cfg['trange'][0] * 1000, hyper_time[0]], 'k',
                lw=6 * cfg['lineweight'])  # 直线部分
        ax.plot(hyper_trace, hyper_time, 'k', lw=4 * cfg['lineweight'])  # 曲线部分
        ax.plot([0, 20], [hyper_time[-1], hyper_time[-1]], 'k', lw=4 * cfg['lineweight'])  # 直线部分
        ax.plot([0, 0], [hyper_time[-1], cfg['trange'][1] * 1000], 'k', lw=4 * cfg['lineweight'])  # 直线部分
        # 第二条曲线
        ax.plot([hyper_trace[0] + trace_num * 20, hyper_trace[0] + trace_num * 20],
                [cfg['trange'][0] * 1000, hyper_time[0]], 'k', lw=6 * cfg['lineweight'])  # 直线部分
        ax.plot(hyper_trace + trace_num * 20, hyper_time, 'k', lw=4 * cfg['lineweight'])  # 曲线部分
        ax.plot([trace_num * 20, trace_num * 20 + 20], [hyper_time[-1], hyper_time[-1]], 'k',
                lw=4 * cfg['lineweight'])  # 直线部分
        ax.plot([trace_num * 20, trace_num * 20], [hyper_time[-1], cfg['trange'][1] * 1000], 'k',
                lw=4 * cfg['lineweight'])  # 直线部分
    else:  # 向右偏
        # 第一条曲线
        ax.plot([0, 0], [cfg['trange'][0] * 1000, hyper_time[0]], 'k',
                lw=6 * cfg['lineweight'])  # 直线部分
        ax.plot(hyper_trace - trace_num * 20, hyper_time, 'k', lw=4 * cfg['lineweight'])  # 曲线部分
        ax.plot([hyper_trace[-1] - trace_num * 20, hyper_trace[-1] - trace_num * 20 + 20],
                [hyper_time[-1], hyper_time[-1]], 'k', lw=4 * cfg['lineweight'])  # 直线部分
        ax.plot([hyper_trace[-1] - trace_num * 20 + 20, hyper_trace[-1] - trace_num * 20 + 20],
                [hyper_time[-1], cfg['trange'][1] * 1000], 'k', lw=4 * cfg['lineweight'])  # 直线部分
        # 第二条曲线
        ax.plot([hyper_trace[0], hyper_trace[0]], [cfg['trange'][0] * 1000, hyper_time[0]], 'k',
                lw=6 * cfg['lineweight'])  # 直线部分
        ax.plot(hyper_trace, hyper_time, 'k', lw=4 * cfg['lineweight'])  # 曲线部分
        ax.plot([hyper_trace[-1], hyper_trace[-1] + 20],
                [hyper_time[-1], hyper_time[-1]], 'k', lw=4 * cfg['lineweight'])  # 直线部分
        ax.plot([hyper_trace[-1] + 20, hyper_trace[-1] + 20],
                [hyper_time[-1], cfg['trange'][1] * 1000], 'k', lw=4 * cfg['lineweight'])  # 直线部分

    # 绘制分层
    if cfg["fcx"] == 1:
        for ii in range(0, len(fctime)):
            if fctime[ii] > fbmax or fctime[ii] < fbmin:
                continue
            ax.plot([0, (stack["trace_num_offset"] + 1) * 20], [fctime[ii], fctime[ii]], 'r', lw=2.5 * cfg['lineweight'])

    # 设置坐标轴
    ax.set_xlim(0, (stack["trace_num_offset"] + 1) * 20)
    # 坐标轴翻转
    # ax.invert_yaxis()
    # ax.invert_xaxis()

    # 隐藏y轴
    # 隐藏 Y 轴
    ax.spines['left'].set_visible(False)  # 隐藏 Y 轴的边框
    ax.spines['right'].set_visible(False)
    ax.yaxis.set_visible(False)  # 隐藏 Y 轴的刻度和标签

    return ax

# 用于地面数据
def wiggle_plot2(file, cfg, ax, gain, fctime, fbmax, fbmin, times=0, stt=0):
    """
    用于绘制地面数据
    :param file: 地面数据
    :param cfg: 参数对象
    :param ax: 坐标系对象
    :param gain: 振幅大小
    :param fctime: 分层深度对应的初至数组
    :param fbmax: 初至最大值
    :param fbmin: 初至最小值
    :param times: 时移，默认为0
    :param stt: 索引，地面数据，井口与检波点距离最近的的道的索引
    :return: 坐标系对象
    """
    perc = 99.0
    data = file["data"]
    rgba = list(cfg['colour']) + [cfg['opacity']]
    sc = np.percentile(data, perc)  # Normalization factor，计算数据集中特定百分位数的值，相当于去除异常数据后的最大值
    stp = int((cfg['trange'][0] * 1000.0 + times) / file["dt"])  # 计算时间点索引，开始，加时移times
    etp = int((cfg['trange'][1] * 1000.0 + times) / file["dt"])  # 计算时间点索引，结束，加时移times

    wigdata = data[stt : stt+cfg['sein'] : cfg["skip"], stp : etp : cfg["skip2"]]  # 对道和时间进行抽稀，并对道和时间进行截取

    xpos = np.arange(20, 20 * cfg['sein'] + 20, 20)  # 生成[20, sein*20]的数组，步长为20。默认每道的深度为20

    for x, trace in zip(xpos, wigdata):
        # Compute high resolution trace.
        # 将道数据与深度结合
        amp = x + gain * 20 * trace / sc

        # 对道数据进行插值
        t = np.arange(stp * file["dt"], etp * file["dt"], file["dt"] * cfg["skip2"])
        # 相当于将数组timerArr的颗粒度由1变为0.1
        hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
        # 利用下行差值公式得到新的amp数组
        hyperamp = np.interp(hypertime, t, amp)

        # 绘制该道数据
        ax.plot(hyperamp, hypertime, 'k', lw=1.2 * cfg['lineweight'])
        # 对波形进行填充
        ax.fill_betweenx(hypertime, hyperamp, x,
                         where=hyperamp > x,
                         facecolor=rgba,
                         lw=0,
                         )

    # 绘制分层
    if cfg["fcx"] == 1:
        for ii in range(0, len(fctime)):
            if fctime[ii] > fbmax or fctime[ii] < fbmin:
                continue
            ax.plot([0, 500], [fctime[ii] + times,
                               fctime[ii] + times], 'r', lw=2.5 * cfg['lineweight'])

    ax.set_xlim(0, 20 * cfg['sein'] + 20)
    ax.invert_yaxis()

    return ax
